import os
import numpy as np
import pandas as pd
import pyBigWig
from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
import random
from tqdm import tqdm
import json
from utils import config

class GenomicDataGenerator:
    def __init__(self, output_dir=None):
        """
        基因组数据生成器
        
        Args:
            output_dir: 输出目录，如果为None则使用配置文件中的目录
        """
        self.config = config['data_generation']
        self.output_dir = output_dir or config['data_dir']
        os.makedirs(self.output_dir, exist_ok=True)
        
        # 设置随机种子
        np.random.seed(config['data_params']['random_state'])
        random.seed(config['data_params']['random_state'])
    
    def generate_genome(self):
        """生成基因组序列"""
        print("生成基因组序列...")
        records = []
        for i in tqdm(range(self.config['n_chromosomes']), desc="生成染色体"):
            # 生成随机DNA序列
            sequence = ''.join(np.random.choice(['A', 'T', 'C', 'G'], size=self.config['genome_size']))
            
            # 创建序列记录
            record = SeqRecord(
                Seq(sequence),
                id=f"chr{i+1}",
                description=f"Simulated chromosome {i+1}"
            )
            records.append(record)
        
        # 保存到FASTA文件
        output_file = os.path.join(self.output_dir, "genome.fa")
        SeqIO.write(records, output_file, "fasta")
        return output_file
    
    def generate_regulatory_regions(self):
        """生成调控区域（增强子和启动子）"""
        print("生成调控区域...")
        
        # 生成增强子
        enhancers = []
        for i in tqdm(range(self.config['n_enhancers']), desc="生成增强子"):
            chrom = f"chr{random.randint(1, self.config['n_chromosomes'])}"
            start = random.randint(0, self.config['genome_size'] - self.config['enhancer_length']['max'])
            end = start + random.randint(
                self.config['enhancer_length']['min'],
                self.config['enhancer_length']['max']
            )
            score = random.random()
            enhancers.append({
                'chrom': chrom,
                'start': start,
                'end': end,
                'name': f'enhancer_{i}',
                'score': score,
                'strand': random.choice(['+', '-'])
            })
        
        # 生成启动子
        promoters = []
        for i in tqdm(range(self.config['n_promoters']), desc="生成启动子"):
            chrom = f"chr{random.randint(1, self.config['n_chromosomes'])}"
            start = random.randint(0, self.config['genome_size'] - self.config['promoter_length']['max'])
            end = start + random.randint(
                self.config['promoter_length']['min'],
                self.config['promoter_length']['max']
            )
            score = random.random()
            promoters.append({
                'chrom': chrom,
                'start': start,
                'end': end,
                'name': f'promoter_{i}',
                'score': score,
                'strand': random.choice(['+', '-'])
            })
        
        # 保存到BED文件
        enhancer_df = pd.DataFrame(enhancers)
        promoter_df = pd.DataFrame(promoters)
        
        enhancer_file = os.path.join(self.output_dir, "enhancers.bed")
        promoter_file = os.path.join(self.output_dir, "promoters.bed")
        
        enhancer_df.to_csv(enhancer_file, sep='\t', index=False, header=False)
        promoter_df.to_csv(promoter_file, sep='\t', index=False, header=False)
        
        return enhancer_file, promoter_file
    
    def generate_bigwig(self):
        """生成组蛋白修饰信号数据"""
        print("生成组蛋白修饰信号...")
        
        def create_signal(filename, params):
            bw = pyBigWig.open(filename, "w")
            
            # 添加染色体头信息
            chroms = [(f"chr{i+1}", self.config['genome_size']) for i in range(self.config['n_chromosomes'])]
            bw.addHeader(chroms)
            
            # 为每条染色体生成信号
            for chrom, size in tqdm(chroms, desc=f"生成{filename}的信号"):
                # 生成基础信号
                signal = np.random.normal(params['base_level'], params['noise_level'], size)
                
                # 添加一些峰
                n_peaks = random.randint(50, 200)
                for _ in range(n_peaks):
                    peak_pos = random.randint(0, size-1)
                    peak_width = random.randint(50, 500)
                    peak_height = random.uniform(
                        params['peak_height']['min'],
                        params['peak_height']['max']
                    )
                    
                    start = max(0, peak_pos - peak_width//2)
                    end = min(size, peak_pos + peak_width//2)
                    
                    # 创建高斯峰
                    x = np.arange(start, end)
                    peak = peak_height * np.exp(-(x - peak_pos)**2 / (2 * (peak_width/4)**2))
                    signal[start:end] += peak[:end-start]
                
                # 确保信号非负
                signal = np.maximum(signal, 0)
                
                # 添加到bigwig文件
                bw.addEntries(
                    chrom,
                    0,
                    values=signal.tolist(),
                    span=1,
                    step=1
                )
            
            bw.close()
        
        # 生成H3K27ac和H3K4me3信号
        h3k27ac_file = os.path.join(self.output_dir, "H3K27ac.bw")
        h3k4me3_file = os.path.join(self.output_dir, "H3K4me3.bw")
        
        create_signal(h3k27ac_file, self.config['histone_modification']['h3k27ac'])
        create_signal(h3k4me3_file, self.config['histone_modification']['h3k4me3'])
        
        return h3k27ac_file, h3k4me3_file
    
    def generate_gene_expression(self):
        """生成基因表达数据"""
        print("生成基因表达数据...")
        
        genes = []
        for i in tqdm(range(self.config['n_genes']), desc="生成基因表达数据"):
            chrom = f"chr{random.randint(1, self.config['n_chromosomes'])}"
            position = random.randint(0, self.config['genome_size'])
            expression = np.random.gamma(2, 2)  # 使用gamma分布模拟表达值
            
            genes.append({
                'gene_id': f'gene_{i}',
                'chrom': chrom,
                'position': position,
                'expression': expression
            })
        
        gene_df = pd.DataFrame(genes)
        output_file = os.path.join(self.output_dir, "gene_expression.csv")
        gene_df.to_csv(output_file, index=False)
        
        return output_file
    
    def generate_all(self):
        """生成所有数据"""
        print(f"开始生成数据到目录: {self.output_dir}")
        
        genome_file = self.generate_genome()
        enhancer_file, promoter_file = self.generate_regulatory_regions()
        h3k27ac_file, h3k4me3_file = self.generate_bigwig()
        gene_expression_file = self.generate_gene_expression()
        
        print("\n生成的文件:")
        print(f"基因组序列: {genome_file}")
        print(f"增强子区域: {enhancer_file}")
        print(f"启动子区域: {promoter_file}")
        print(f"H3K27ac信号: {h3k27ac_file}")
        print(f"H3K4me3信号: {h3k4me3_file}")
        print(f"基因表达: {gene_expression_file}")
        
        # 生成数据统计信息
        stats = {
            'genome_size': self.config['genome_size'],
            'n_chromosomes': self.config['n_chromosomes'],
            'n_enhancers': self.config['n_enhancers'],
            'n_promoters': self.config['n_promoters'],
            'n_genes': self.config['n_genes'],
            'enhancer_length': self.config['enhancer_length'],
            'promoter_length': self.config['promoter_length'],
            'histone_modification': self.config['histone_modification']
        }
        
        # 保存统计信息
        stats_file = os.path.join(self.output_dir, "data_stats.json")
        with open(stats_file, 'w') as f:
            json.dump(stats, f, indent=4)
        
        print(f"\n数据统计信息已保存到: {stats_file}")
        return stats

def main():
    # 创建生成器并生成数据
    generator = GenomicDataGenerator()
    stats = generator.generate_all()

if __name__ == "__main__":
    main() 